#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project : aimix
@File    : auto_synthesis.py
@IDE     : PyCharm
@Author  : admin
@Date    : 2025/4/10 13:15
"""
import random
from copy import deepcopy
from concurrent.futures import ThreadPoolExecutor, as_completed

from core.video.base import BaseVideo
from utils.load_file import load_files
from core.video.processor import load_video


class AutoSynthesis(BaseVideo):
    """自动合成"""
    def __init__(self, video_path, num=2, rule="分镜数量", times=1, select_rule="random", *args, **kwargs):
        super(AutoSynthesis, self).__init__(video_path, *args, **kwargs)
        self.num = int(num)  # 镜头数量
        self.rule = rule  # 分配规则
        self.times = int(times)  # 镜头时长
        self.select_rule = select_rule  # 选择规则
        self.video_groups = self.load_video_groups(video_path)  # 视频组路径


    def load_video_groups(self, video_path):
        with ThreadPoolExecutor(max_workers=4) as executor:
            tasks = load_files(video_path)
            res = []
            features = []
            for task in tasks:
                features.append(executor.submit(load_video, task))
            for future in as_completed(features):
                res.append(future.result())
        return res


    def distritbute_tasks(self):
        """
        分配任务
        :param task_files: 任务文件列表
        :param times: 镜头数量
        :param count: 任务数量
        :param ruler: 任务分配规则
        :return: 任务列表
        """
        video_groups = deepcopy(self.video_groups)
        if not video_groups or self.times <= 0 or self.num_outputs <= 0:
            return []
        results = []
        if self.rule == "分镜数量":
            task_files = deepcopy(self.video_groups)
            if self.select_rule == "random":
                for i in range(self.num_outputs):
                    results.append(random.sample(task_files, self.num))
            elif self.select_rule == "unique":
                for i in range(self.num_outputs):
                    res = []
                    while len(res) < self.num and len(task_files):
                        res.append(task_files.pop(random.randint(0, len(task_files) - 1)))
                    if len(res) == self.num:
                        results.append(res)
        elif self.rule == "视频时长":
            if self.select_rule == "unique":
                # 所有视频只能用一次
                available_videos = deepcopy(video_groups)
                random.shuffle(available_videos)
                while available_videos and len(results) < self.num_outputs:
                    group = []
                    total_duration = 0

                    i = 0
                    while i < len(available_videos):
                        video = available_videos[i]
                        group.append(video)
                        total_duration += video["duration"]
                        i += 1
                        if total_duration >= self.times:
                            break

                    if total_duration >= self.times:
                        results.append(group)
                        # 移除已使用的视频
                        for vid in group:
                            available_videos.remove(vid)
                    else:
                        break  # 无法再拼出一组满足时长的视频

            else:  # "random"
                for _ in range(self.num_outputs):
                    pool = deepcopy(video_groups)
                    random.shuffle(pool)
                    group = []
                    used_paths = set()
                    total_duration = 0

                    for video in pool:
                        if video["filename"] in used_paths:
                            continue  # 避免组内重复
                        group.append(video)
                        used_paths.add(video["filename"])
                        total_duration += video["duration"]
                        if total_duration >= self.times:
                            break
                    if total_duration >= self.times:
                        results.append(group)
        else:
            raise ValueError("rule must be '分镜数量' or '视频时长'")

        return results


    def generate_combinations(self):
        files_item = self.distritbute_tasks()
        return [[i.get("filename") for i in j] for j in files_item]

